from matrixslow.core import Variable
from matrixslow.ops import MatMul, Add, Logistic, ReLU


def fc(input, input_size, size, activation):
    '''
    全连接层
    :param input: 节点对象，保存输入向量
    :param input_size: 整数，输入向量的维数
    :param size: 整数，该层的神经元数量，也是输出向量的维数
    :param activation: 字符串，激活函数的类型
    :return:
    '''
    weights = Variable((size, input_size), init=True, trainable=True)   # 保存权值矩阵的变量节点
    bias = Variable((size, 1), init=True, trainable=True)               # 保存偏置向量的变量节点
    affine = Add(MatMul(weights, input), bias)                          # Add类节点，将权重与输入相乘然后再加上bias节点得到的（仿射变换/线性变换）

    if activation == 'ReLU':
        return ReLU(affine)
    elif activation == 'Logistic':
        return Logistic(affine)
    else:
        return affine